from typing import List, Optional, Dict, Any
from sqlalchemy.orm import Session
from app.models.task import Task, TaskStatus
from app.schemas.task import TaskCreate, TaskUpdate
from app.core.celery_config import celery_app
import asyncio


class TaskManager:
    """任务管理器，负责任务的创建、调度和状态管理"""
    
    async def create_task(self, db: Session, task: TaskCreate) -> Task:
        """创建新任务"""
        # 创建任务记录
        db_task = Task(
            name=task.name,
            description=task.description,
            engine_name=task.engine_name,
            target_id=task.target_id,
            config=task.config,
            status=TaskStatus.PENDING,
            progress=0.0
        )
        
        db.add(db_task)
        db.commit()
        db.refresh(db_task)
        
        return db_task
    
    async def get_tasks(self, db: Session, skip: int = 0, limit: int = 100, **filters) -> List[Task]:
        """获取任务列表"""
        query = db.query(Task)
        
        # 应用过滤条件
        for key, value in filters.items():
            if hasattr(Task, key):
                query = query.filter(getattr(Task, key) == value)
        
        return query.offset(skip).limit(limit).all()
    
    async def get_task(self, db: Session, task_id: int) -> Optional[Task]:
        """根据ID获取任务"""
        return db.query(Task).filter(Task.id == task_id).first()
    
    async def update_task(self, db: Session, task_id: int, task_update: TaskUpdate) -> Optional[Task]:
        """更新任务信息"""
        db_task = await self.get_task(db, task_id)
        if not db_task:
            return None
        
        # 只更新提供的字段
        update_data = task_update.model_dump(exclude_unset=True)
        for field, value in update_data.items():
            setattr(db_task, field, value)
        
        db.commit()
        db.refresh(db_task)
        
        return db_task
    
    async def delete_task(self, db: Session, task_id: int) -> bool:
        """删除任务"""
        db_task = await self.get_task(db, task_id)
        if not db_task:
            return False
        
        # 如果任务正在运行，先停止
        if db_task.status == TaskStatus.RUNNING and db_task.celery_task_id:
            celery_app.control.revoke(db_task.celery_task_id, terminate=True)
        
        db.delete(db_task)
        db.commit()
        
        return True
    
    async def start_task(self, db: Session, task_id: int) -> Task:
        """开始任务"""
        db_task = await self.get_task(db, task_id)
        if not db_task:
            raise ValueError(f"Task {task_id} not found")
        
        if db_task.status == TaskStatus.RUNNING:
            raise ValueError("Task is already running")
        
        # 更新任务状态
        db_task.status = TaskStatus.RUNNING
        db_task.progress = 0.0
        
        # 提交到Celery执行
        from app.task_manager.tasks import run_fuzz_task
        result = run_fuzz_task.delay(task_id)
        db_task.celery_task_id = result.id
        
        db.commit()
        db.refresh(db_task)
        
        return db_task
    
    async def stop_task(self, db: Session, task_id: int) -> Task:
        """停止任务"""
        db_task = await self.get_task(db, task_id)
        if not db_task:
            raise ValueError(f"Task {task_id} not found")
        
        if db_task.status != TaskStatus.RUNNING:
            raise ValueError("Task is not running")
        
        # 取消Celery任务
        if db_task.celery_task_id:
            celery_app.control.revoke(db_task.celery_task_id, terminate=True)
        
        # 更新任务状态
        db_task.status = TaskStatus.FAILED
        
        db.commit()
        db.refresh(db_task)
        
        return db_task
    
    async def pause_task(self, db: Session, task_id: int) -> Task:
        """暂停任务"""
        db_task = await self.get_task(db, task_id)
        if not db_task:
            raise ValueError(f"Task {task_id} not found")
        
        if db_task.status != TaskStatus.RUNNING:
            raise ValueError("Task is not running")
        
        # 更新任务状态
        db_task.status = TaskStatus.PAUSED
        
        # 暂停Celery任务（实际上可能需要特殊处理）
        # 这里简化处理，实际应用中可能需要更复杂的暂停机制
        
        db.commit()
        db.refresh(db_task)
        
        return db_task
    
    async def resume_task(self, db: Session, task_id: int) -> Task:
        """恢复任务"""
        db_task = await self.get_task(db, task_id)
        if not db_task:
            raise ValueError(f"Task {task_id} not found")
        
        if db_task.status != TaskStatus.PAUSED:
            raise ValueError("Task is not paused")
        
        # 更新任务状态
        db_task.status = TaskStatus.RUNNING
        
        # 如果之前的任务已经结束，创建新的任务
        if db_task.celery_task_id:
            from celery.result import AsyncResult
            result = AsyncResult(db_task.celery_task_id, app=celery_app)
            if result.ready():
                from app.task_manager.tasks import run_fuzz_task
                new_result = run_fuzz_task.delay(task_id)
                db_task.celery_task_id = new_result.id
        
        db.commit()
        db.refresh(db_task)
        
        return db_task